import collections
import copy
import datetime
import itertools
import random
import time
import traceback

from absl import app, flags

import numpy

from coin.strategy.kr_hard_arb.arb_info import create_arb_info_general
from coin.strategy.kr_hard_arb.util import *

FLAGS = flags.FLAGS

flags.DEFINE_string('currency', '', '')

flags.DEFINE_string('exchanges', '', '')

flags.DEFINE_string('key_files', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_bool('dry_run', False, 'dry_run')


def get_ref_price_mean(exchs, currency):
  prices = []
  for exch in exchs:
    prices.append(exch.feed.get_best_price())
    if prices[-1] is None:
      raise ValueError((exch.exchange, exch.feed.get_best_price()))
  return numpy.mean([price.mid_price for price in prices])


def get_ref_price(info, currency):
  return get_ref_price_mean(info.exchs, currency)


def has_enough_bal(bal, threshold):
  return bal.available >= threshold * 1.01


def buy_exch1_sell_exch2(exch2_order,
                         exch1_order,
                         buy_price,
                         buy_qty,
                         sell_price,
                         sell_qty,
                         currency,
                         qty_threshold):
  if FLAGS.dry_run:
    return
  jobs = []
  if buy_qty > qty_threshold:
    jobs.append(
        ThreadRunner(exch1_order.query_submit,
                     exch1_order.get_product(currency),
                     "bid",
                     buy_price,
                     buy_qty))
  if sell_qty > qty_threshold:
    jobs.append(
        ThreadRunner(exch2_order.query_submit,
                     exch2_order.get_product(currency),
                     "ask",
                     sell_price,
                     sell_qty))
  try:
    execute_thread_runners(jobs, "Buy exch1 / Sell exch2")
    sleep_seconds(2)
  except Exception:
    traceback.print_exc()
  execute_thread_runners([
      ThreadRunner(exch1_order.query_cancel_product, exch1_order.get_product(currency)),
      ThreadRunner(exch2_order.query_cancel_product, exch2_order.get_product(currency))
  ],
                         "Cancel All")


def buy_exch2_sell_exch1(exch2_order,
                         exch1_order,
                         buy_price,
                         buy_qty,
                         sell_price,
                         sell_qty,
                         currency,
                         qty_threshold):
  if FLAGS.dry_run:
    return
  jobs = []
  if buy_qty > qty_threshold:
    jobs.append(
        ThreadRunner(exch2_order.query_submit,
                     exch2_order.get_product(currency),
                     "bid",
                     buy_price,
                     buy_qty))
  if sell_qty > qty_threshold:
    jobs.append(
        ThreadRunner(exch1_order.query_submit,
                     exch1_order.get_product(currency),
                     "ask",
                     sell_price,
                     sell_qty))
  try:
    execute_thread_runners(jobs, "Buy exch2 / Sell exch1")
    sleep_seconds(2)
  except Exception:
    traceback.print_exc()
  execute_thread_runners([
      ThreadRunner(exch1_order.query_cancel_product, exch1_order.get_product(currency)),
      ThreadRunner(exch2_order.query_cancel_product, exch2_order.get_product(currency))
  ],
                         "Cancel All")


def truncate(value):
  return int(value * 1000) / 1000.0


def get_currency_balance_impl(order, currency):
  balance = order.get_balance_cached()
  if balance is None or balance.msg is None:
    raise RuntimeError("balance.msg is None")

  for each_balance in balance.msg.each_balance:
    if each_balance.currency == currency or (each_balance.currency == "BCC"
                                             and currency.lower() == "bch"):
      return each_balance

  from coin.exchange.upbit_v1.kr_rest.order_client import UpbitV1Order

  if isinstance(order, UpbitV1Order):
    from coin.proto.coin_query_pb2 import CurrencyBalance
    currency_balance = CurrencyBalance()
    currency_balance.currency = currency
    currency_balance.currency_native = currency
    currency_balance.available = 0.0
    currency_balance.hold = 0.0
    currency_balance.total = 0.0
    return currency_balance
  else:
    raise ValueError(currency, balance.msg)
  return None


def get_currency_balance(info, exch, currency):
  balance_ret = get_currency_balance_impl(exch.order, currency)
  balance = copy.deepcopy(balance_ret)
  offset = info.config.get("offsets", {}).get(currency, {}).get(exch.exchange, None)
  if offset:
    balance.total -= offset
  return balance


def get_currency_balance_map(info, currency):
  currency_bal = collections.OrderedDict()
  for exch in info.exchs:
    currency_bal[exch.exchange] = get_currency_balance(info, exch, currency).total
  return currency_bal


def print_status(info, currency):
  ccy_price = get_ref_price(info, currency)

  krws = get_currency_balance_map(info, 'KRW')
  ccys = get_currency_balance_map(info, currency)

  total_ccy = sum(ccys.values())
  total_krw = sum(krws.values())
  total_value = sum(krws.values()) + ccy_price * total_ccy
  print("Estimated Value: %s" % ("{:,.0f}".format(total_value)))
  print("Krw Value: %s, %s" % ("{:,.0f}".format(total_krw), datetime.datetime.now()))
  print("%s Value: %s * %s = %s" % (currency, ccy_price, total_ccy, ccy_price * total_ccy))

  for ccy, format_str, total, ccymap in [('KRW', "%s: %9.0f", total_krw, krws),
                                         (currency, "%s: %9.4f", total_ccy, ccys)]:
    print(
        "%s \t %s" %
        (ccy,
         "\t".join([format_str % ("Total", total)]
                   + [format_str % (exchange, ccy_value)
                      for exchange, ccy_value in ccymap.items()])))


def get_buy_sell_qty(info, currency):
  config = info.config
  lot_size_krw = config["lot_size_krw"]

  total_ccy = sum(get_currency_balance_map(info, currency).values())
  ref_price = get_ref_price(info, currency)

  lot_size = (lot_size_krw / float(ref_price)) * (1 + random.random() * 0.2)
  base_qty = max(lot_size - abs(config["qty"] - total_ccy), 0)
  buy_qty = truncate(base_qty + max(config["qty"] - total_ccy, 0))
  sell_qty = truncate(base_qty + max(total_ccy - config["qty"], 0))
  buy_qty_lot_adj = buy_qty - min(buy_qty, lot_size)
  sell_qty_lot_adj = sell_qty - min(sell_qty, lot_size)
  adj_qty_max = max(buy_qty_lot_adj, sell_qty_lot_adj)
  buy_qty -= adj_qty_max
  sell_qty -= adj_qty_max
  buy_qty = max(buy_qty, 0)
  sell_qty = max(sell_qty, 0)
  assert buy_qty >= 0
  assert sell_qty >= 0
  return (buy_qty, sell_qty)


def is_expected_deposit(info, currency):
  total_coin = sum(get_currency_balance_map(info, currency).values())
  expected_total_coin = info.config["qty"]
  qty_offset_ratio = info.config.get("qty_offset_ratio", 0.20)
  diff = abs(expected_total_coin - total_coin)
  print("Total: %s, Expected: %s" % (total_coin, expected_total_coin))
  print("Diff:", diff, "Threshold:", expected_total_coin * qty_offset_ratio)
  if diff > expected_total_coin * qty_offset_ratio:
    return False
  else:
    return True


die_flag = False


def trade_coin(info, currency, verbose):
  if not is_expected_deposit(info, currency):
    print()
    print(("It seems that %s deposit is quite different "
           "from expectation. Skip trading...") % currency)
    return

  exch_sent_order = set()
  exch_price_cache = {}
  exch_balance_cache = {}

  for exch1, exch2 in itertools.combinations(info.exchs, 2):
    # prevent from ordering subsequently
    if exch1.exchange in exch_sent_order or exch2.exchange in exch_sent_order:
      continue
    if (hasattr(exch1.feed, "exchange_time") and exch1.feed.exchange_time):
      ref_time1 = exch1.feed.exchange_time
    else:
      ref_time1 = exch1.feed.fetched_time

    if (hasattr(exch2.feed, "exchange_time") and exch2.feed.exchange_time):
      ref_time2 = exch2.feed.exchange_time
    else:
      ref_time2 = exch2.feed.fetched_time

    diff_time = abs(ref_time1 - ref_time2)
    if (diff_time > 5 * 1e+9):
      print("tdiff=%.3f, %s time1 = %s, %s time2 = %s" %
            (diff_time / 1e9, exch1.exchange, ref_time1, exch2.exchange, ref_time2))
      if diff_time > 30 * 1e+9:
        die_flag = True
        raise ValueError("diff is serious")
      return
    elif verbose:
      print("tdiff=%.3f, normal" % (diff_time / 1e9))

    price_det_notional = info.config["price_det_notional"]
    rate_threshold_pre = info.config["rate_threshold"]
    urgent_rate_threshold_pre = info.config["urgent_rate_threshold"]
    qty_threshold = info.config["min_qty"]
    max_qty = info.config["max_qty"]

    if exch1.exchange not in exch_price_cache:
      exch_price_cache[exch1.exchange] = exch1.feed.get_price_by_notional(price_det_notional)
    if exch2.exchange not in exch_price_cache:
      exch_price_cache[exch2.exchange] = exch2.feed.get_price_by_notional(price_det_notional)
    exch1_price = exch_price_cache[exch1.exchange]
    exch2_price = exch_price_cache[exch2.exchange]

    print("%25s: %s" % (exch1.exchange, exch1_price))
    print("%25s: %s" % (exch2.exchange, exch2_price))

    if exch1_price is None or exch2_price is None:
      return

    edge1 = info.config["edge"].get(exch1.exchange, 0)
    edge2 = info.config["edge"].get(exch2.exchange, 0)

    rate_threshold = rate_threshold_pre + (edge1 + edge2)
    urgent_rate_threshold = urgent_rate_threshold_pre + (edge1 + edge2)

    for exch in [exch1, exch2]:
      for query_ccy in ('KRW', currency):
        if (exch.exchange, query_ccy) not in exch_balance_cache:
          exch_balance_cache[(exch.exchange,
                              query_ccy)] = get_currency_balance(info, exch, query_ccy)

    exch2_krw = exch_balance_cache[(exch2.exchange, 'KRW')]
    exch1_krw = exch_balance_cache[(exch1.exchange, 'KRW')]
    exch2_bal = exch_balance_cache[(exch2.exchange, currency)]
    exch1_bal = exch_balance_cache[(exch1.exchange, currency)]

    buy_qty, sell_qty = get_buy_sell_qty(info, currency)

    to_exch2 = float(exch2_price.bid_price) / float(exch1_price.ask_price) - 1
    to_exch1 = float(exch1_price.bid_price) / float(exch2_price.ask_price) - 1

    urgent_threshold = (info.config["qty"] * 0.4) / len(info.exchs)

    max_coin1 = info.config.get("max_coin", {}).get(currency, {}).get(exch1.exchange, None)
    max_coin2 = info.config.get("max_coin", {}).get(currency, {}).get(exch2.exchange, None)

    if max_coin1 is None:
      max_coin1 = info.config["qty"]

    if max_coin2 is None:
      max_coin2 = info.config["qty"]

    urgent_threshold1 = max_coin1 * 0.2
    urgent_threshold2 = max_coin2 * 0.2

    if urgent_threshold1 is None:
      urgent_threshold1 = urgent_threshold
    else:
      urgent_threshold1 = min(urgent_threshold, urgent_threshold1)

    if urgent_threshold2 is None:
      urgent_threshold2 = urgent_threshold
    else:
      urgent_threshold2 = min(urgent_threshold, urgent_threshold2)

    lack1 = exch1_bal.total < urgent_threshold1
    lack2 = exch2_bal.total < urgent_threshold2

    can_buy1 = exch1_bal.total < max_coin1
    can_buy2 = exch2_bal.total < max_coin2

    urgency12 = (lack1 and not lack2)
    urgency21 = (not lack1 and lack2)

    threshold12 = urgent_rate_threshold if urgency12 else rate_threshold
    threshold21 = urgent_rate_threshold if urgency21 else rate_threshold

    if verbose:
      print("%-25s Threshold: %8.4f+(%8.4f+%8.4f) Urgent: %8.4f+(%8.4f+%8.4f)" %
            ("%s-%s" % (exch1.exchange, exch2.exchange),
             rate_threshold,
             rate_threshold_pre, (edge1 + edge2),
             urgent_rate_threshold,
             urgent_rate_threshold_pre, (edge1 + edge2)))
      print("lack    %10s: %8.2f < %8.2f ? %5s"
            "%10s: %8.2f < %8.2f ? %5s" % (exch1.exchange,
                                           exch1_bal.total,
                                           urgent_threshold1,
                                           lack1,
                                           exch2.exchange,
                                           exch2_bal.total,
                                           urgent_threshold2,
                                           lack2))
      print("can_buy %10s: %8.2f < %8.2f ? %5s"
            "%10s: %8.2f < %8.2f ? %5s" % (exch1.exchange,
                                           exch1_bal.total,
                                           max_coin1,
                                           can_buy1,
                                           exch2.exchange,
                                           exch2_bal.total,
                                           max_coin2,
                                           can_buy2))
      print("%25s: %8.4f %25s: %8.4f " %
            ("%s Edge" % exch1.exchange, edge1, "%s Edge" % exch2.exchange, edge2))
      print("%25s: %8s %25s: %8s" % ("-%s+%s Urgent" % (exch1.exchange, exch2.exchange),
                                     urgency21,
                                     "-%s+%s Urgent" % (exch2.exchange, exch1.exchange),
                                     urgency12))
      print("%25s: %8.4f %25s: %8.4f" % ("-%s+%s Thres" % (exch1.exchange, exch2.exchange),
                                         threshold21,
                                         "-%s+%s Thres" % (exch2.exchange, exch1.exchange),
                                         threshold12))
    print("%25s: %8.4f %25s: %8.4f" % ("-%s+%s" % (exch1.exchange, exch2.exchange),
                                       to_exch1,
                                       "-%s+%s" % (exch2.exchange, exch1.exchange),
                                       to_exch2))

    # buy exchange 1
    if to_exch2 > threshold12 and can_buy1:
      buy_qty_max = (max_coin1 - exch1_bal.total) * 0.99
      sell_qty_max = exch2_bal.available * 0.99
      buy_adj_qty = max(buy_qty - buy_qty_max, 0)
      sell_adj_qty = max(sell_qty - sell_qty_max, 0)
      adj_qty = max(buy_adj_qty, sell_adj_qty)
      buy_qty = max(truncate(buy_qty - adj_qty), 0)
      sell_qty = max(truncate(sell_qty - adj_qty), 0)
      assert buy_qty - adj_qty >= 0, (exch1_bal,
                                      exch2_bal,
                                      max_coin1,
                                      buy_qty,
                                      sell_qty,
                                      buy_adj_qty,
                                      sell_adj_qty)
      assert sell_qty - adj_qty >= 0, (exch1_bal,
                                       exch2_bal,
                                       max_coin1,
                                       buy_qty,
                                       sell_qty,
                                       buy_adj_qty,
                                       sell_adj_qty)

      print("Buy %s / Sell %s" % (exch1.exchange, exch2.exchange))
      print("Urgency: %s, lack1: %s, lack2: %s" % (urgency12, lack1, lack2))
      print("Rate: %.4f%%" % (to_exch2 * 100))
      print("Buy  %.4f at %d" % (buy_qty, exch1_price.ask_price))
      print("Sell %.4f at %d" % (sell_qty, exch2_price.bid_price))
      if has_enough_bal(exch1_krw, buy_qty * exch1_price.ask_price) and \
         has_enough_bal(exch2_bal, sell_qty) and \
         (buy_qty > qty_threshold or sell_qty > qty_threshold) and \
         (buy_qty < max_qty and sell_qty < max_qty):
        exch_sent_order.add(exch1.exchange)
        exch_sent_order.add(exch2.exchange)
        buy_exch1_sell_exch2(exch2.order,
                             exch1.order,
                             exch1_price.ask_price,
                             buy_qty,
                             exch2_price.bid_price,
                             sell_qty,
                             currency,
                             qty_threshold)
      else:
        print(max_coin1, buy_qty, sell_qty, buy_adj_qty, sell_adj_qty)
        print("Not enough Balance")

    # buy exchange 2
    if to_exch1 > threshold21 and can_buy2:
      buy_qty_max = (max_coin2 - exch2_bal.total) * 0.99
      sell_qty_max = exch1_bal.available * 0.99
      buy_adj_qty = max(buy_qty - buy_qty_max, 0)
      sell_adj_qty = max(sell_qty - sell_qty_max, 0)
      adj_qty = max(buy_adj_qty, sell_adj_qty)
      buy_qty = max(truncate(buy_qty - adj_qty), 0)
      sell_qty = max(truncate(sell_qty - adj_qty), 0)
      assert buy_qty - adj_qty >= 0, (exch1_bal,
                                      exch2_bal,
                                      max_coin2,
                                      buy_qty,
                                      sell_qty,
                                      buy_adj_qty,
                                      sell_adj_qty)
      assert sell_qty - adj_qty >= 0, (exch1_bal,
                                       exch2_bal,
                                       max_coin2,
                                       buy_qty,
                                       sell_qty,
                                       buy_adj_qty,
                                       sell_adj_qty)

      print("Buy %s / Sell %s" % (exch2.exchange, exch1.exchange))
      print("Urgency: %s, lack1: %s, lack2: %s" % (urgency21, lack1, lack2))
      print("Rate: %.4f%%" % (to_exch1 * 100))
      print("Buy  %.4f at %d" % (buy_qty, exch2_price.ask_price))
      print("Sell %.4f at %d" % (sell_qty, exch1_price.bid_price))
      if has_enough_bal(exch2_krw, buy_qty * exch2_price.ask_price) and \
         has_enough_bal(exch1_bal, sell_qty) and \
         (buy_qty > qty_threshold or sell_qty > qty_threshold) and \
         (buy_qty < max_qty and sell_qty < max_qty):
        exch_sent_order.add(exch1.exchange)
        exch_sent_order.add(exch2.exchange)
        buy_exch2_sell_exch1(exch2.order,
                             exch1.order,
                             exch2_price.ask_price,
                             buy_qty,
                             exch1_price.bid_price,
                             sell_qty,
                             currency,
                             qty_threshold)
      else:
        print(max_coin1, buy_qty, sell_qty, buy_adj_qty, sell_adj_qty)
        print("Not enough Balance")


def trade(info, currency, verbose):
  try:
    print()
    execute_thread_runners([ThreadRunner(exch.feed.query_book) for exch in info.exchs]
                           + [ThreadRunner(exch.order.query_balance) for exch in info.exchs],
                           "Collecting %s Data" % currency)
    print_status(info, currency)
    trade_coin(info, currency, verbose)
  except KeyboardInterrupt:
    raise
  except Exception as exc:
    print("A problem detected. Try again")
    print(exc)
    traceback.print_exc()


def main(_):
  currency = FLAGS.currency

  info = create_arb_info_general(currency,
                                 FLAGS.exchanges.split(","),
                                 FLAGS.key_files.split(","),
                                 FLAGS.config_filename)

  verbose_every_n = 10
  verbose_count = 0

  while (True):
    time.sleep(1)
    print()
    print(datetime.datetime.now())
    print("Checking...")
    trade(info, currency, (verbose_count % verbose_every_n) == 0)
    verbose_count += 1
    if die_flag:
      break
  return 0


if __name__ == '__main__':
  app.run(main)
